package com.boardour.toupin.push;

import android.hardware.display.DisplayManager;
import android.hardware.display.VirtualDisplay;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.projection.MediaProjection;
import android.util.Log;
import android.view.Surface;

import com.boardour.socket.client.WifiClientPresenter;
import com.boardour.toupin.VideoUtils;
import com.boardour.wifip2p.receiver.AddressManager;

import java.io.Closeable;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

public class PushEncode implements Runnable, Closeable {
    //private static final String MIME_TYPE = MediaFormat.MIMETYPE_VIDEO_MPEG4;
    private static final String MIME_TYPE = VideoUtils.MIME_TYPE;
    private AtomicBoolean isStart = new AtomicBoolean(false);
    private AtomicBoolean isClose = new AtomicBoolean(false);
    private MediaProjection mediaProjection;
    private MediaFormat format;
    private MediaCodec mediaCodec;
    private Surface surface;
    private VirtualDisplay virtualDisplay;
    private byte[] SPS_PPS;
    public static int NAL_I = 5; // H264
    public static int NAL_SPS = 7;// H264
    private WifiClientPresenter client;

    public PushEncode(MediaProjection mediaProjection, WifiClientPresenter client) {
        this.client = client;
        this.mediaProjection = mediaProjection;
        initMediaCode();
    }

    private void initMediaCode() {
        format = MediaFormat.createVideoFormat(MIME_TYPE, VideoUtils.DECODE_WIDTH, VideoUtils.DECODE_HEIGHT);
        format.setInteger(MediaFormat.KEY_BIT_RATE, VideoUtils.BIT_RATE);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, VideoUtils.FRAME_RATE);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, VideoUtils.I_FRAME_INTERVAL);
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
        try {
            mediaCodec = MediaCodec.createEncoderByType(MIME_TYPE);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    public void start() {
        if (mediaCodec == null) {
            return;
        }
        if (isStart.compareAndSet(false, true)) {
            Executors.newSingleThreadExecutor().execute(this);
        }
    }

    private void release() {
        mediaCodec.stop();
        mediaCodec.release();
        virtualDisplay.release();
        mediaProjection.stop();
    }

    @Override
    public void close() {
        if (isClose.compareAndSet(false, true)) {
            isStart.set(false);
            client.close();
        }
    }

    @Override
    public void run() {
        Log.e("TAG", "start...");
        mediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        surface = mediaCodec.createInputSurface();
        virtualDisplay = mediaProjection.createVirtualDisplay(
                "TouPin", VideoUtils.DECODE_WIDTH, VideoUtils.DECODE_HEIGHT, 1,
                DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC,
                surface, null, null);
        mediaCodec.start();
        MediaCodec.BufferInfo info = new MediaCodec.BufferInfo();
        while (isStart.get()) {
            int index = mediaCodec.dequeueOutputBuffer(info, 40);
            if (index == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                // csd-0
                MediaFormat outputFormat = mediaCodec.getOutputFormat();
                ByteBuffer csd0 = outputFormat.getByteBuffer("csd-0");
                ByteBuffer csd1 = outputFormat.getByteBuffer("csd-1");
                if (csd1 != null) {
                    // H264
                    SPS_PPS = new byte[csd0.remaining() + csd1.remaining()];
                    csd0.get(SPS_PPS, 0, csd0.remaining());
                    csd1.get(SPS_PPS, csd0.remaining() + 1, csd1.remaining());
                    Log.e("TAG", "SPS PPS");
                } else if (csd0 != null) {
                    // MP4 H265 都走这个
                    SPS_PPS = new byte[csd0.remaining()];
                    csd0.get(SPS_PPS, 0, csd0.remaining());
                    Log.e("TAG", "VPS SPS PPS");
                }
            } else if (index >= 0) {
                ByteBuffer buffer = mediaCodec.getOutputBuffer(index);
                if (VideoUtils.MIME_TYPE == MediaFormat.MIMETYPE_VIDEO_MPEG4) {
                    dealFrameMp4(info, buffer);
                } else {
                    dealFrameH264(info, buffer);
                }
                mediaCodec.releaseOutputBuffer(index, false);
            }
        }
        Log.e("TAG", "start end ...");
        release();
    }

    private void dealFrameMp4(MediaCodec.BufferInfo info, ByteBuffer buffer) {
        boolean keyFrame = (info.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0;
        //Log.e("TAG", "keyFrame:" + keyFrame+":"+(info.flags ));
        int offset = 8;
        /*if (buffer.get(2) == 0x01) {
            offset = 9;
        }*/
        int type = buffer.get(offset) & 0xc0;
        /*if (type == 0x00) {
            // I Frame
            Log.e("TAG", "NAL_I:" + type);
        }else if (type == 0x40) {
            Log.e("TAG", "NAL_P:" + type);
        } else if (type == 0x80){
            Log.e("TAG", "NAL_B:" + type);
        }*/

        final byte[] bytes = new byte[info.size];
        buffer.get(bytes);
        byte[] data = new byte[SPS_PPS.length + info.size];
        System.arraycopy(SPS_PPS, 0, data, 0, SPS_PPS.length);
        System.arraycopy(bytes, 0, data, SPS_PPS.length, bytes.length);
        Log.e("CODE", "NAL_I:" + data.length);
        if (data.length < 100) {
            return;
        }
        callback(data, info.presentationTimeUs);

        /*if (keyFrame) {

            final byte[] bytes = new byte[info.size];
            buffer.get(bytes);
            byte[] data = new byte[SPS_PPS.length + info.size];
            System.arraycopy(SPS_PPS, 0, data, 0, SPS_PPS.length);
            System.arraycopy(bytes, 0, data, SPS_PPS.length, bytes.length);
            Log.e("CODE", "NAL_I:" + data.length);
            if (data.length < 100) {
                return;
            }
            callback(data, info.presentationTimeUs);
        } else {
            final byte[] data = new byte[info.size];
            buffer.get(data);
            Log.e("CODE", "NAL_B:" + data.length);
            callback(data, info.presentationTimeUs);
        }*/


        /*int type = buffer.get(8) & 0xc0;
        if (type == 0x00) {
            // I Frame
            Log.e("TAG", "NAL_I:" + type);
            final byte[] bytes = new byte[info.size];
            buffer.get(bytes);
            byte[] data = new byte[SPS_PPS.length + info.size];
            System.arraycopy(SPS_PPS, 0, data, 0, SPS_PPS.length);
            System.arraycopy(bytes, 0, data, SPS_PPS.length, bytes.length);
            callback(data, info.presentationTimeUs);
        } else if (type == 0x40 || type == 0x80) {
            if (type == 0x40) {
                Log.e("TAG", "NAL_P:" + type);
            } else {
                Log.e("TAG", "NAL_B:" + type);
            }
            final byte[] data = new byte[info.size];
            buffer.get(data);
            callback(data, info.presentationTimeUs);
        }*/
    }


    private void dealFrameH264(MediaCodec.BufferInfo info, ByteBuffer buffer) {
        int offset = 4;
        if (buffer.get(2) == 0x01) {
            offset = 3;
        }
        // int type = (buffer.get(offset) & 0x7E) >> 1; // H265
        int type = buffer.get(offset) & 0x1F; // h264
        if (type == NAL_SPS) {
            SPS_PPS = new byte[info.size];
            buffer.get(SPS_PPS);
            //Log.e("TAG", "NAL_SPS");
        } else if (type == NAL_I) {
            final byte[] bytes = new byte[info.size];
            buffer.get(bytes);
            byte[] data = new byte[SPS_PPS.length + info.size];
            //Log.e("TAG", "NAL_I:" + data.length+":"+SPS_PPS.length);
            System.arraycopy(SPS_PPS, 0, data, 0, SPS_PPS.length);
            System.arraycopy(bytes, 0, data, SPS_PPS.length, bytes.length);
            callback(data, info.presentationTimeUs);
        } else {
            // Log.e("TAG", "NAL_Other");
            final byte[] data = new byte[info.size];
            buffer.get(data);
            callback(data, info.presentationTimeUs);
        }
    }

    private void callback(byte[] data, long time) {
        AddressManager manager = AddressManager.getInstance();
        InetAddress address = manager.getAddress();
        if (address == null) {
            return;
        }
        if (!manager.isGroupOwner()) {
            realSend(address, data, time);
        } else {
            String ip = client.getIp();
            if (ip != null) {
                realSend(ip, data, time);
            }
        }
    }

    private void realSend(String ip, byte[] data, long time) {
        client.connectServer(ip);
        client.send(0, data, time);
    }

    private void realSend(InetAddress address, byte[] data, long time) {
        client.connectServer(address);
        client.send(0, data, time);
    }

    public void stop() {

    }
}
